import json
import torch
from torch.utils.data import Dataset
from modelscope import AutoTokenizer
from dataset import utils

## 打开json格式的微调训练
with open("dataset/chatGLM3_dataFormatted_sample.json", "r", encoding="UTF-8") as j_file:
    j_dict = json.load(j_file)

## 定义模型路径
model_dir = "C:\\Users\\16014\\.cache\\modelscope\\hub\\models\\ZhipuAI\\chatglm3-6b"
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)


# 定义一个名为ChatDataset的类，该类继承自Dataset类，用于处理聊天数据集
class ChatDataset(Dataset):
    # 初始化函数，接收以下参数：
    # conversations: 一个包含对话数据的字典，默认为j_dict
    # tokenizer: 用于tokenization的工具，默认为tokenizer
    # max_tokens: 最大token数量限制，默认为None
    def __init__(self, conversations: {} = j_dict, tokenizer=tokenizer, max_tokens=None):
        # 通过super()调用父类Dataset的初始化函数
        super(ChatDataset, self).__init__()

        # 使用utils.preprocess函数预处理对话数据，得到处理后的数据字典
        data_dict = utils.preprocess(conversations, tokenizer, max_tokens)

        # 从处理后的数据字典中提取input_ids和labels，并保存到类的属性中
        self.input_ids = data_dict["input_ids"]
        self.labels = data_dict["labels"]

        # 重写__len__方法，返回处理后的input_ids的长度，即数据集的大小

    def __len__(self):
        return len(self.input_ids)

        # 重写__getitem__方法，使得可以通过索引i获取数据集中第i个样本的input_ids和labels

    def __getitem__(self, i):
        return dict(input_ids=self.input_ids[i], labels=self.labels[i])


# 定义一个名为 DataCollatorForChatDataset 的类，它继承自 object 类。
class DataCollatorForChatDataset(object):
    """
    Collate examples for supervised fine-tuning.
    """
    # 初始化函数，这里没有接收特定的参数。
    def __init__(self):
        # 初始化一个属性 padding_value，并设置其值为0。这个属性后续用于 padding 操作。
        self.padding_value = 0

        # 定义一个特殊方法 __call__，这使得类的实例能够像函数一样被调用。

    def __call__(self, instances):
        # instances 参数应该是一个包含多个实例的列表，每个实例都是一个字典，包含 'input_ids' 和 'labels'。
        # 通过列表推导式，分别提取每个实例的 'input_ids' 和 'labels'，并组成新的列表。
        input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "labels"))

        # 使用 torch.nn.utils.rnn.pad_sequence 函数对 input_ids 列表进行 padding 操作，使得所有的序列长度一致。
        # 参数 batch_first=True 表示输入的数据是 batch-major，即第一个维度是 batch 维度。
        # 参数 padding_value=self.padding_value 表示用 0 进行 padding。
        input_ids = torch.nn.utils.rnn.pad_sequence(input_ids, batch_first=True, padding_value=self.padding_value)

        # 与上面类似，对 labels 列表进行 padding 操作，但是这里使用 -100 进行 padding。
        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=-100)

        # 返回一个字典，包含经过处理后的 input_ids, labels, 以及根据 input_ids 生成的 attention_mask。
        # attention_mask 是一个布尔类型的张量，它的作用是在模型处理输入时，告诉模型哪些部分是真正的内容，哪些部分是 padding。
        return dict(
            input_ids=input_ids,
            labels=labels,
            attention_mask=input_ids.ne(self.padding_value),
        )


if __name__ == '__main__':
    ChatDataset()
    print(tokenizer.pad_token_id)